[Kernel][AMD] Avoid D2H copy and cumsum kernel #22683
Conversation
Current AITER attention for prefill incurs one D2H copy (which causes synchonrization between CPU and GPU) and one cumsum kernel. The D2H copy is used to get the actual KV tokens on the host side for intermediate buffer allocation; this value should not change during the inference, so we can pre-compute it during AttentionMetadata construction time. And similarly, the cumsum kernel can be done in AttentionMedata construction time Signed-off-by: Xiaozhu <[email protected]>
mxz297
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to optimize the AITER attention backend for prefill on ROCm by avoiding a device-to-host copy and a redundant cumsum kernel execution. The changes achieve this by pre-computing cu_seq_lens and num_actual_kv_tokens during the AttentionMetadata construction, which is a sound optimization strategy.
My review includes a suggestion to correct a type hint for the newly added cu_seq_lens attribute.
Additionally, I've identified a potential critical issue in the AiterFlashAttentionImpl.forward method. It appears that for prefill requests, both the prefill attention kernel (flash_attn_varlen_func) and the decode attention kernel (paged_attention_v1) are called sequentially. This would cause the output of the prefill kernel to be overwritten by the decode kernel, likely resulting in incorrect outputs and performance degradation. I strongly recommend wrapping the decode kernel call in an else block to ensure it's only executed for decode requests. While this issue pre-exists this PR, it is critical to fix for the backend to function correctly.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Michael Goin <[email protected]>
github-actions
bot
added
the
ready
|
@HAIAI @mgoin First PR to vllm oss, so not quite sure about the CIs results. It looks to me tpu tests are not related? the v1-test failure says but seems like it is on NV platform, for which my pr should have no impact... Can you give some advice on the CI failures? |
|
Hi @mxz297 unfortunately the CI has lots of failures at the moment, we usually try to keep it in a better state. Please ignore these for now and I will request force merge. |
Signed-off-by: Xiaozhu <[email protected]> Signed-off-by: Michael Goin <[email protected]> Co-authored-by: Michael Goin <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Diego-Castan <[email protected]>
Signed-off-by: Xiaozhu <[email protected]> Signed-off-by: Michael Goin <[email protected]> Co-authored-by: Michael Goin <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: Xiaozhu <[email protected]> Signed-off-by: Michael Goin <[email protected]> Co-authored-by: Michael Goin <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
The PR breaks full graph mode. |
|
@fsx950223 could you file an issue for it please? |
Signed-off-by: Xiaozhu <[email protected]> Signed-off-by: Michael Goin <[email protected]> Co-authored-by: Michael Goin <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: Xiaozhu <[email protected]> Signed-off-by: Michael Goin <[email protected]> Co-authored-by: Michael Goin <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Xiao Yu <[email protected]>
Signed-off-by: Xiaozhu <[email protected]> Signed-off-by: Michael Goin <[email protected]> Co-authored-by: Michael Goin <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
6 participants
Current AITER attention for prefill incurs one D2H copy (which causes synchonrization between CPU and GPU) and one cumsum kernel.
The D2H copy is used to get the actual KV tokens on the host side for intermediate buffer allocation; this value should not change during the inference, so we can pre-compute it during AttentionMetadata construction time. And similarly, the cumsum kernel can be done in AttentionMedata construction time
We can see there is memcpy and gpu bubbles after the change.
Saw about 10% prefill latency improvement
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
Test Plan
Test Result
(Optional) Documentation Update